import cv2
import asyncio
import websockets
import json
import base64
import time
import logging
import argparse
import threading
from datetime import datetime
import socket

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class VideoCapture:
    def __init__(self, server_host='xxx.xxx.xxx.xxx', server_port=8765, camera_id=0, 
                 fps=30, quality=85, resize_width=640):
        self.server_host = server_host
        self.server_port = server_port
        self.camera_id = camera_id
        self.fps = fps
        self.quality = quality
        self.resize_width = resize_width
        
        self.cap = None
        self.ws = None
        self.is_running = False
        self.is_connected = False
        self.is_capturing = False  # 默认关闭采集，等待服务端指令
        self.frame_count = 0
        self.start_time = time.time()
        
        # 状态监控
        self.last_frame_time = 0
        self.connection_attempts = 0
        self.max_reconnect_attempts = 5
        
        # 自动重连相关
        self.auto_reconnect_enabled = True
        self.reconnect_interval = 5  # 5秒重连间隔
        self.reconnect_task = None
        
    def initialize_camera(self):
        """初始化摄像头"""
        try:
            # 尝试打开摄像头
            self.cap = cv2.VideoCapture(self.camera_id)
            
            if not self.cap.isOpened():
                # 如果数字ID失败，尝试字符串路径
                if isinstance(self.camera_id, int):
                    logger.warning(f"Camera {self.camera_id} not found, trying /dev/video{self.camera_id}")
                    self.cap = cv2.VideoCapture(f"/dev/video{self.camera_id}")
                
                if not self.cap.isOpened():
                    raise Exception(f"Cannot open camera {self.camera_id}")
            
            # 设置摄像头参数
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.resize_width)
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, int(self.resize_width * 3/4))
            self.cap.set(cv2.CAP_PROP_FPS, self.fps)
            
            # 获取实际参数
            actual_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            actual_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            actual_fps = self.cap.get(cv2.CAP_PROP_FPS)
            
            logger.info(f"Camera initialized: {actual_width}x{actual_height} @ {actual_fps}fps")
            logger.info("摄像头已就绪，等待服务端采集指令...")
            return True
            
        except Exception as e:
            logger.error(f"Failed to initialize camera: {e}")
            return False
    
    def capture_frame(self):
        """捕获单帧图像"""
        if not self.cap or not self.cap.isOpened():
            return None
            
        ret, frame = self.cap.read()
        if not ret:
            logger.warning("Failed to capture frame")
            return None
        
        try:
            # 调整图像尺寸
            height, width = frame.shape[:2]
            if width != self.resize_width:
                new_height = int(height * self.resize_width / width)
                frame = cv2.resize(frame, (self.resize_width, new_height))
            
            # 编码为JPEG
            encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), self.quality]
            result, encoded_img = cv2.imencode('.jpg', frame, encode_param)
            
            if result:
                # 转换为base64
                img_base64 = base64.b64encode(encoded_img).decode('utf-8')
                return img_base64
            else:
                logger.warning("Failed to encode frame")
                return None
                
        except Exception as e:
            logger.error(f"Error processing frame: {e}")
            return None
    
    async def connect_to_server(self, is_reconnect=False):
        """连接到服务器"""
        max_retries = 3 if is_reconnect else 10
        retry_delay = 2 if is_reconnect else 5
        
        for attempt in range(max_retries):
            if not self.is_running:  # 如果程序已停止，不再尝试连接
                return False               
            try:
                ws_url = f"ws://{self.server_host}:{self.server_port}"
                connect_type = "重连" if is_reconnect else "连接"
                logger.info(f"{connect_type}到服务器: {ws_url} (第 {attempt + 1}/{max_retries} 次尝试)")
                
                self.ws = await websockets.connect(ws_url, ping_interval=20, ping_timeout=10)
                
                # 注册为摄像头
                await self.ws.send(json.dumps({
                    'type': 'camera_register',
                    'camera_id': f"camera_{self.camera_id}",
                    'timestamp': time.time()
                }))
                
                # 等待响应
                response = await asyncio.wait_for(self.ws.recv(), timeout=5.0)
                response_data = json.loads(response)
                
                if response_data.get('success'):
                    success_msg = "重连成功" if is_reconnect else "连接成功"
                    logger.info(f"{success_msg}，已向服务器注册")
                    logger.info("已连接到服务端，等待采集指令...")
                    self.is_connected = True
                    self.connection_attempts = 0
                    
                    # 停止自动重连任务
                    if self.reconnect_task and not self.reconnect_task.done():
                        self.reconnect_task.cancel()
                        self.reconnect_task = None
                    
                    return True
                else:
                    logger.error(f"注册失败: {response_data.get('message')}")
                    
            except asyncio.TimeoutError:
                logger.error("连接超时")
            except websockets.exceptions.ConnectionClosed:
                logger.error("服务器关闭连接")
            except Exception as e:
                logger.error(f"连接失败: {e}")
            
            if attempt < max_retries - 1:
                wait_msg = f"等待 {retry_delay} 秒后重试..."
                logger.info(wait_msg)
                await asyncio.sleep(retry_delay)
        
        return False
    
    async def start_auto_reconnect(self):
        """启动自动重连任务"""
        if not self.auto_reconnect_enabled or not self.is_running:
            return
            
        if self.reconnect_task and not self.reconnect_task.done():
            return  # 重连任务已在运行
        
        logger.info(f"启动自动重连任务，每 {self.reconnect_interval} 秒尝试重连...")
        self.reconnect_task = asyncio.create_task(self._auto_reconnect_loop())
    
    async def _auto_reconnect_loop(self):
        """自动重连循环"""
        while self.is_running and self.auto_reconnect_enabled:
            if not self.is_connected:
                logger.info(f"尝试自动重连到服务器...")
                success = await self.connect_to_server(is_reconnect=True)
                
                if success:
                    logger.info("自动重连成功！")
                    # 重连成功后启动消息处理
                    asyncio.create_task(self.message_handler())
                    break
                else:
                    logger.warning(f"自动重连失败，{self.reconnect_interval} 秒后再次尝试...")
            else:
                break  # 已连接，退出重连循环
            
            # 等待重连间隔
            try:
                await asyncio.sleep(self.reconnect_interval)
            except asyncio.CancelledError:
                logger.info("自动重连任务被取消")
                break
    
    async def send_frame(self, frame_data):
        """发送视频帧到服务器"""
        if not self.ws or not self.is_connected or not self.is_capturing:
            return False
            
        try:
            message = json.dumps({
                'type': 'frame_data',
                'camera_id': f"camera_{self.camera_id}",
                'frame': frame_data,
                'timestamp': time.time(),
                'frame_count': self.frame_count
            })
            
            await self.ws.send(message)
            return True
            
        except websockets.exceptions.ConnectionClosed:
            logger.warning("发送帧时连接断开")
            self.is_connected = False
            # 启动自动重连
            await self.start_auto_reconnect()
            return False
        except Exception as e:
            logger.error(f"发送帧错误: {e}")
            self.is_connected = False
            # 启动自动重连
            await self.start_auto_reconnect()
            return False
    
    async def send_control_response(self, command, success=True, message=""):
        """发送控制命令响应"""
        if not self.ws or not self.is_connected:
            return
            
        try:
            response = json.dumps({
                'type': 'control_response',
                'command': command,
                'success': success,
                'message': message,
                'camera_id': f"camera_{self.camera_id}",
                'timestamp': time.time()
            })
            await self.ws.send(response)
        except websockets.exceptions.ConnectionClosed:
            logger.warning("发送控制响应时连接断开")
            self.is_connected = False
            await self.start_auto_reconnect()
        except Exception as e:
            logger.error(f"发送控制响应错误: {e}")
            self.is_connected = False
            await self.start_auto_reconnect()
    
    async def handle_control_command(self, command):
        """处理控制命令"""
        try:
            if command == 'start_capture':
                if not self.is_capturing:
                    self.is_capturing = True
                    self.frame_count = 0  # 重置帧计数
                    self.start_time = time.time()  # 重置计时器
                    logger.info("收到开始采集指令，开始视频采集...")
                    await self.send_control_response(command, True, "视频采集已开始")
                else:
                    logger.info("收到开始采集指令，但采集已在运行")
                    await self.send_control_response(command, True, "视频采集已在运行")
                    
            elif command == 'stop_capture':
                if self.is_capturing:
                    self.is_capturing = False
                    logger.info("收到停止采集指令，停止视频采集")
                    await self.send_control_response(command, True, "视频采集已停止")
                else:
                    logger.info("收到停止采集指令，但采集已停止")
                    await self.send_control_response(command, True, "视频采集已停止")
                    
            else:
                logger.warning(f"未知命令: {command}")
                await self.send_control_response(command, False, f"未知命令: {command}")
                
        except Exception as e:
            logger.error(f"处理控制命令错误: {e}")
            await self.send_control_response(command, False, f"执行命令失败: {str(e)}")
    
    async def message_handler(self):
        """处理来自服务器的消息"""
        try:
            async for message in self.ws:
                try:
                    data = json.loads(message)
                    
                    if data.get('type') == 'control_command':
                        command = data.get('command')
                        if command:
                            logger.info(f"收到控制命令: {command}")
                            await self.handle_control_command(command)
                    
                    elif data.get('type') == 'register_response':
                        if data.get('success'):
                            logger.info("服务器确认注册成功")
                        else:
                            logger.error(f"注册失败: {data.get('message')}")
                    
                except json.JSONDecodeError:
                    logger.error("从服务器收到无效JSON数据")
                except Exception as e:
                    logger.error(f"处理服务器消息错误: {e}")
                    
        except websockets.exceptions.ConnectionClosed:
            logger.warning("消息处理器: 连接断开")
            self.is_connected = False
            # 启动自动重连
            if self.is_running:
                await self.start_auto_reconnect()
        except Exception as e:
            logger.error(f"消息处理器错误: {e}")
            self.is_connected = False
            # 启动自动重连
            if self.is_running:
                await self.start_auto_reconnect()
    
    async def video_stream_loop(self):
        """视频流发送循环"""
        frame_interval = 1.0 / self.fps
        
        while self.is_running:
            loop_start = time.time()
            
            # 只有在连接状态且采集开启时才捕获和发送帧
            if self.is_connected and self.is_capturing:
                # 捕获帧
                frame_data = self.capture_frame()
                if frame_data:
                    # 发送帧
                    success = await self.send_frame(frame_data)
                    if success:
                        self.frame_count += 1
                        self.last_frame_time = time.time()
                        
                        # 每100帧输出一次统计信息
                        if self.frame_count % 100 == 0:
                            elapsed = time.time() - self.start_time
                            avg_fps = self.frame_count / elapsed
                            logger.info(f"已发送 {self.frame_count} 帧, 平均FPS: {avg_fps:.2f}")
                
                # 控制帧率
                loop_time = time.time() - loop_start
                sleep_time = max(0, frame_interval - loop_time)
                if sleep_time > 0:
                    await asyncio.sleep(sleep_time)
            else:
                # 如果未连接或不采集，稍微等待一下避免CPU占用过高
                await asyncio.sleep(0.1)
    
    async def reconnect(self):
        """手动重新连接到服务器（已弃用，由自动重连替代）"""
        logger.info("手动重连已被自动重连功能替代")
        await self.start_auto_reconnect()
    
    def display_preview(self):
        """显示预览窗口（调试用）"""
        if not self.cap or not self.cap.isOpened():
            return
            
        while self.is_running:
            ret, frame = self.cap.read()
            if ret:
                # 调整尺寸用于显示
                height, width = frame.shape[:2]
                if width > 800:  # 限制预览窗口大小
                    new_width = 800
                    new_height = int(height * new_width / width)
                    display_frame = cv2.resize(frame, (new_width, new_height))
                else:
                    display_frame = frame
                
                # 添加状态信息
                connection_status = "已连接" if self.is_connected else "断开连接"
                capture_status = "采集中" if self.is_capturing else "等待指令"
                status_text = f"FPS: {self.fps} | 帧数: {self.frame_count} | {connection_status} | {capture_status}"
                cv2.putText(display_frame, status_text, (10, 30), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
                
                # 显示连接状态
                if not self.is_connected:
                    cv2.putText(display_frame, "CONNECTION LOST - AUTO RECONNECTING", 
                               (display_frame.shape[1]//2 - 250, display_frame.shape[0]//2 - 20), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
                
                # 显示采集状态
                if not self.is_capturing and self.is_connected:
                    cv2.putText(display_frame, "WAITING FOR START COMMAND", 
                               (display_frame.shape[1]//2 - 200, display_frame.shape[0]//2 + 20), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)
                elif self.is_capturing and self.is_connected:
                    cv2.putText(display_frame, "CAPTURING", 
                               (display_frame.shape[1]//2 - 80, display_frame.shape[0]//2 + 20), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
                
                cv2.imshow(f'Camera {self.camera_id} - Preview', display_frame)
                
                # 按'q'退出
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    logger.info("预览窗口被用户关闭")
                    self.stop()
                    break
            else:
                time.sleep(0.1)
    
    async def start_async(self, show_preview=False):
        """异步启动采集"""
        self.is_running = True
        
        # 初始化摄像头
        if not self.initialize_camera():
            logger.error("摄像头初始化失败，程序退出...")
            return
        
        # 首次连接到服务器
        if not await self.connect_to_server():
            logger.error("初始连接服务器失败，启动自动重连...")
            # 即使首次连接失败，也启动自动重连
            await self.start_auto_reconnect()
        
        # 启动预览窗口（在单独线程中）
        if show_preview:
            preview_thread = threading.Thread(target=self.display_preview)
            preview_thread.daemon = True
            preview_thread.start()
        
        try:
            # 创建任务列表
            tasks = [self.video_stream_loop()]
            
            # 如果初始连接成功，启动消息处理
            if self.is_connected:
                tasks.append(self.message_handler())
            
            # 同时运行所有任务
            await asyncio.gather(*tasks, return_exceptions=True)
            
        except KeyboardInterrupt:
            logger.info("收到键盘中断信号")
        except Exception as e:
            logger.error(f"主循环错误: {e}")
        finally:
            self.cleanup()
    
    def start(self, show_preview=False):
        """启动采集（同步版本）"""
        try:
            asyncio.run(self.start_async(show_preview))
        except KeyboardInterrupt:
            logger.info("程序被用户中断")
        except Exception as e:
            logger.error(f"程序错误: {e}")
    
    def stop(self):
        """停止采集"""
        logger.info("正在停止视频采集...")
        self.is_running = False
        self.auto_reconnect_enabled = False
        
        # 取消自动重连任务
        if self.reconnect_task and not self.reconnect_task.done():
            self.reconnect_task.cancel()
    
    def cleanup(self):
        """清理资源"""
        logger.info("正在清理资源...")
        
        self.auto_reconnect_enabled = False
        
        # 取消自动重连任务
        if self.reconnect_task and not self.reconnect_task.done():
            self.reconnect_task.cancel()
        
        # 关闭WebSocket连接
        if self.ws:
            try:
                asyncio.create_task(self.ws.close())
            except:
                pass
        
        # 释放摄像头
        if self.cap:
            self.cap.release()
        
        cv2.destroyAllWindows()
        logger.info("资源清理完成")

def get_available_cameras():
    """获取可用的摄像头列表"""
    available_cameras = []
    
    # 测试前10个摄像头索引
    for i in range(10):
        cap = cv2.VideoCapture(i)
        if cap.isOpened():
            ret, _ = cap.read()
            if ret:
                available_cameras.append(i)
                logger.info(f"发现摄像头: 索引 {i}")
            cap.release()
        else:
            # 在Linux系统上尝试/dev/video*设备
            try:
                cap = cv2.VideoCapture(f"/dev/video{i}")
                if cap.isOpened():
                    ret, _ = cap.read()
                    if ret:
                        available_cameras.append(f"/dev/video{i}")
                        logger.info(f"发现摄像头: /dev/video{i}")
                    cap.release()
            except:
                pass
    
    return available_cameras

def test_connection(host, port):
    """测试服务器连接"""
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(5)
        result = sock.connect_ex((host, port))
        sock.close()
        return result == 0
    except Exception as e:
        logger.error(f"连接测试失败: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='视频监控系统 - 采集端 (带自动重连功能)')
    parser.add_argument('--server', '-s', default='xxx.xxx.xxx.xxx', 
                       help='服务器地址 (默认: xxx.xxx.xxx.xxx)')
    parser.add_argument('--port', '-p', type=int, default=8765, 
                       help='服务器端口 (默认: 8765)')
    parser.add_argument('--camera', '-c', default=0, 
                       help='摄像头ID (默认: 0)')
    parser.add_argument('--fps', '-f', type=int, default=20, 
                       help='帧率 (默认: 20)')
    parser.add_argument('--quality', '-q', type=int, default=70, 
                       help='JPEG质量 1-100 (默认: 70)')
    parser.add_argument('--width', '-w', type=int, default=640, 
                       help='视频宽度 (默认: 640)')
    parser.add_argument('--preview', action='store_true', 
                       help='显示预览窗口')
    parser.add_argument('--list-cameras', action='store_true', 
                       help='列出可用摄像头')
    
    args = parser.parse_args()
    
    # 列出可用摄像头
    if args.list_cameras:
        logger.info("正在扫描可用摄像头...")
        cameras = get_available_cameras()
        if cameras:
            logger.info("可用摄像头:")
            for cam in cameras:
                logger.info(f"  - {cam}")
        else:
            logger.warning("未找到可用摄像头")
        return
    
    # 转换摄像头ID
    try:
        camera_id = int(args.camera)
    except ValueError:
        camera_id = args.camera
    
    # 测试服务器连接
    logger.info(f"测试服务器连接: {args.server}:{args.port}")
    if not test_connection(args.server, args.port):
        logger.warning(f"无法连接到服务器 {args.server}:{args.port}")
        logger.warning("程序将启动并自动尝试重连...")
    
    # 创建并启动采集端
    logger.info("启动视频采集端 (带自动重连功能)...")
    logger.info(f"配置: 服务器={args.server}:{args.port}, 摄像头={camera_id}")
    logger.info(f"参数: FPS={args.fps}, 质量={args.quality}, 宽度={args.width}")
    logger.info("功能: 连接断开后每5秒自动重连")
    logger.info("注意: 采集端已启动，但不会自动开始采集，需要等待服务端指令")
    
    capture = VideoCapture(
        server_host=args.server,
        server_port=args.port,
        camera_id=camera_id,
        fps=args.fps,
        quality=args.quality,
        resize_width=args.width
    )
    
    try:
        capture.start(show_preview=args.preview)
    except KeyboardInterrupt:
        logger.info("程序被用户中断")
    except Exception as e:
        logger.error(f"程序异常: {e}")
    finally:
        logger.info("程序退出")

if __name__ == "__main__":
    main()